{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0df23120",
   "metadata": {
    "colab_type": "text",
    "id": "EgiF12Hf1Dhs"
   },
   "source": [
    "This notebook provides examples to go along with the [textbook](http://manipulation.csail.mit.edu/pose.html).  I recommend having both windows open, side-by-side!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f67c1b64",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "eeMrMI0-1Dhu"
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "from copy import copy\n",
    "from enum import Enum\n",
    "\n",
    "import numpy as np\n",
    "from pydrake.all import (\n",
    "    AbstractValue,\n",
    "    AddMultibodyPlantSceneGraph,\n",
    "    Concatenate,\n",
    "    DiagramBuilder,\n",
    "    InputPortIndex,\n",
    "    LeafSystem,\n",
    "    MeshcatVisualizer,\n",
    "    Parser,\n",
    "    PiecewisePolynomial,\n",
    "    PiecewisePose,\n",
    "    PointCloud,\n",
    "    PortSwitch,\n",
    "    RandomGenerator,\n",
    "    RigidTransform,\n",
    "    RollPitchYaw,\n",
    "    Simulator,\n",
    "    StartMeshcat,\n",
    "    UniformlyRandomRotationMatrix,\n",
    ")\n",
    "\n",
    "from manipulation import ConfigureParser, FindResource, running_as_notebook\n",
    "from manipulation.clutter import GenerateAntipodalGraspCandidate\n",
    "from manipulation.meshcat_utils import AddMeshcatTriad\n",
    "from manipulation.pick import (\n",
    "    MakeGripperCommandTrajectory,\n",
    "    MakeGripperFrames,\n",
    "    MakeGripperPoseTrajectory,\n",
    ")\n",
    "from manipulation.scenarios import AddIiwaDifferentialIK, ycb\n",
    "from manipulation.station import (\n",
    "    AddPointClouds,\n",
    "    MakeHardwareStation,\n",
    "    add_directives,\n",
    "    load_scenario,\n",
    ")\n",
    "\n",
    "\n",
    "class NoDiffIKWarnings(logging.Filter):\n",
    "    def filter(self, record):\n",
    "        return not record.getMessage().startswith(\"Differential IK\")\n",
    "\n",
    "\n",
    "logging.getLogger(\"drake\").addFilter(NoDiffIKWarnings())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70fc8875",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db218f7c",
   "metadata": {},
   "source": [
    "# Putting it all together\n",
    "\n",
    "In the code above, we worked with a point cloud using functions.  To assemble this into a full-stack manipulation system, we need to specify the timing semantics of when those functions are called.  That's precisely what Drake's systems framework provides.  I've introduced the systems below:\n",
    "- `GraspSelector` system that takes the camera inputs and outputs the desired grasp,\n",
    "- `Planner` system that takes this pose estimate (and the state of the robot), computes the trajectory, and stores that trajectory in its Context so that it can output the instantaneous command.  This system also keeps track of a `mode` ; the logical state of task execution.\n",
    "\n",
    "I've made only a modest attempt to make this robust.  Enough to show you some of the task-level planning logic, but not enough to work every time.  If things go sideways, then go ahead and hit ESCAPE and rerun the sim.\n",
    "\n",
    "It's fun to think about the remaining limitations.  I would group them into:\n",
    "1. MultibodyPlant can crash; almost always during at time almost zero, if my simple random initial conditions strategy puts objects in deep penetration\n",
    "2. Motion planning. Sometimes it runs into the cameras, knocks the objects into the bins, or gets itself wound up into an unfortunate state for Differential IK.\n",
    "3. Perception. Since there is no notion of object here, our simple grasp selection strategy can pick objects up at the edges or make double-picks.  (Sometimes the gripper gets itself in the point cloud, and the grasp selector picks a candidate grasp on the robot itself, resulting in an apparent airball!  That one is silly and could be fixed with only a little more work)\n",
    "4. More than pick and place. Sometimes objects get stuck in the corners, for instance, and we really need strategies other than pure grasping to get them out.\n",
    "\n",
    "We'll address these pretty soundly in the next few chapters!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bb9d301",
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.default_rng(135)  # this is for python\n",
    "generator = RandomGenerator(rng.integers(0, 1000))  # this is for c++\n",
    "\n",
    "\n",
    "# Another diagram for the objects the robot \"knows about\": gripper, cameras, bins.  Think of this as the model in the robot's head.\n",
    "def make_internal_model():\n",
    "    builder = DiagramBuilder()\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
    "    parser = Parser(plant)\n",
    "    ConfigureParser(parser)\n",
    "    parser.AddModelsFromUrl(\"package://manipulation/clutter_planning.dmd.yaml\")\n",
    "    plant.Finalize()\n",
    "    return builder.Build()\n",
    "\n",
    "\n",
    "# Takes 3 point clouds (in world coordinates) as input, and outputs and estimated pose for the mustard bottle.\n",
    "class GraspSelector(LeafSystem):\n",
    "    def __init__(self, plant, bin_instance, camera_body_indices):\n",
    "        LeafSystem.__init__(self)\n",
    "        model_point_cloud = AbstractValue.Make(PointCloud(0))\n",
    "        self.DeclareAbstractInputPort(\"cloud0_W\", model_point_cloud)\n",
    "        self.DeclareAbstractInputPort(\"cloud1_W\", model_point_cloud)\n",
    "        self.DeclareAbstractInputPort(\"cloud2_W\", model_point_cloud)\n",
    "        self.DeclareAbstractInputPort(\n",
    "            \"body_poses\", AbstractValue.Make([RigidTransform()])\n",
    "        )\n",
    "\n",
    "        port = self.DeclareAbstractOutputPort(\n",
    "            \"grasp_selection\",\n",
    "            lambda: AbstractValue.Make((np.inf, RigidTransform())),\n",
    "            self.SelectGrasp,\n",
    "        )\n",
    "        port.disable_caching_by_default()\n",
    "\n",
    "        # Compute crop box.\n",
    "        context = plant.CreateDefaultContext()\n",
    "        bin_body = plant.GetBodyByName(\"bin_base\", bin_instance)\n",
    "        X_B = plant.EvalBodyPoseInWorld(context, bin_body)\n",
    "        margin = 0.001  # only because simulation is perfect!\n",
    "        a = X_B.multiply(\n",
    "            [-0.22 + 0.025 + margin, -0.29 + 0.025 + margin, 0.015 + margin]\n",
    "        )\n",
    "        b = X_B.multiply([0.22 - 0.1 - margin, 0.29 - 0.025 - margin, 2.0])\n",
    "        self._crop_lower = np.minimum(a, b)\n",
    "        self._crop_upper = np.maximum(a, b)\n",
    "\n",
    "        self._internal_model = make_internal_model()\n",
    "        self._internal_model_context = (\n",
    "            self._internal_model.CreateDefaultContext()\n",
    "        )\n",
    "        self._rng = np.random.default_rng()\n",
    "        self._camera_body_indices = camera_body_indices\n",
    "\n",
    "    def SelectGrasp(self, context, output):\n",
    "        body_poses = self.get_input_port(3).Eval(context)\n",
    "        pcd = []\n",
    "        for i in range(3):\n",
    "            cloud = self.get_input_port(i).Eval(context)\n",
    "            pcd.append(cloud.Crop(self._crop_lower, self._crop_upper))\n",
    "            pcd[i].EstimateNormals(radius=0.1, num_closest=30)\n",
    "\n",
    "            # Flip normals toward camera\n",
    "            X_WC = body_poses[self._camera_body_indices[i]]\n",
    "            pcd[i].FlipNormalsTowardPoint(X_WC.translation())\n",
    "        merged_pcd = Concatenate(pcd)\n",
    "        down_sampled_pcd = merged_pcd.VoxelizedDownSample(voxel_size=0.005)\n",
    "\n",
    "        costs = []\n",
    "        X_Gs = []\n",
    "        # TODO(russt): Take the randomness from an input port, and re-enable\n",
    "        # caching.\n",
    "        for i in range(100 if running_as_notebook else 2):\n",
    "            cost, X_G = GenerateAntipodalGraspCandidate(\n",
    "                self._internal_model,\n",
    "                self._internal_model_context,\n",
    "                down_sampled_pcd,\n",
    "                self._rng,\n",
    "            )\n",
    "            if np.isfinite(cost):\n",
    "                costs.append(cost)\n",
    "                X_Gs.append(X_G)\n",
    "\n",
    "        if len(costs) == 0:\n",
    "            # Didn't find a viable grasp candidate\n",
    "            X_WG = RigidTransform(\n",
    "                RollPitchYaw(-np.pi / 2, 0, np.pi / 2), [0.5, 0, 0.22]\n",
    "            )\n",
    "            output.set_value((np.inf, X_WG))\n",
    "        else:\n",
    "            best = np.argmin(costs)\n",
    "            output.set_value((costs[best], X_Gs[best]))\n",
    "\n",
    "\n",
    "class PlannerState(Enum):\n",
    "    WAIT_FOR_OBJECTS_TO_SETTLE = 1\n",
    "    PICKING_FROM_X_BIN = 2\n",
    "    PICKING_FROM_Y_BIN = 3\n",
    "    GO_HOME = 4\n",
    "\n",
    "\n",
    "class Planner(LeafSystem):\n",
    "    def __init__(self, plant):\n",
    "        LeafSystem.__init__(self)\n",
    "        self._gripper_body_index = plant.GetBodyByName(\"body\").index()\n",
    "        self.DeclareAbstractInputPort(\n",
    "            \"body_poses\", AbstractValue.Make([RigidTransform()])\n",
    "        )\n",
    "        self._x_bin_grasp_index = self.DeclareAbstractInputPort(\n",
    "            \"x_bin_grasp\", AbstractValue.Make((np.inf, RigidTransform()))\n",
    "        ).get_index()\n",
    "        self._y_bin_grasp_index = self.DeclareAbstractInputPort(\n",
    "            \"y_bin_grasp\", AbstractValue.Make((np.inf, RigidTransform()))\n",
    "        ).get_index()\n",
    "        self._wsg_state_index = self.DeclareVectorInputPort(\n",
    "            \"wsg_state\", 2\n",
    "        ).get_index()\n",
    "\n",
    "        self._mode_index = self.DeclareAbstractState(\n",
    "            AbstractValue.Make(PlannerState.WAIT_FOR_OBJECTS_TO_SETTLE)\n",
    "        )\n",
    "        self._traj_X_G_index = self.DeclareAbstractState(\n",
    "            AbstractValue.Make(PiecewisePose())\n",
    "        )\n",
    "        self._traj_wsg_index = self.DeclareAbstractState(\n",
    "            AbstractValue.Make(PiecewisePolynomial())\n",
    "        )\n",
    "        self._times_index = self.DeclareAbstractState(\n",
    "            AbstractValue.Make({\"initial\": 0.0})\n",
    "        )\n",
    "        self._attempts_index = self.DeclareDiscreteState(1)\n",
    "\n",
    "        self.DeclareAbstractOutputPort(\n",
    "            \"X_WG\",\n",
    "            lambda: AbstractValue.Make(RigidTransform()),\n",
    "            self.CalcGripperPose,\n",
    "        )\n",
    "        self.DeclareVectorOutputPort(\"wsg_position\", 1, self.CalcWsgPosition)\n",
    "\n",
    "        # For GoHome mode.\n",
    "        num_positions = 7\n",
    "        self._iiwa_position_index = self.DeclareVectorInputPort(\n",
    "            \"iiwa_position\", num_positions\n",
    "        ).get_index()\n",
    "        self.DeclareAbstractOutputPort(\n",
    "            \"control_mode\",\n",
    "            lambda: AbstractValue.Make(InputPortIndex(0)),\n",
    "            self.CalcControlMode,\n",
    "        )\n",
    "        self.DeclareAbstractOutputPort(\n",
    "            \"reset_diff_ik\",\n",
    "            lambda: AbstractValue.Make(False),\n",
    "            self.CalcDiffIKReset,\n",
    "        )\n",
    "        self._q0_index = self.DeclareDiscreteState(num_positions)  # for q0\n",
    "        self._traj_q_index = self.DeclareAbstractState(\n",
    "            AbstractValue.Make(PiecewisePolynomial())\n",
    "        )\n",
    "        self.DeclareVectorOutputPort(\n",
    "            \"iiwa_position_command\", num_positions, self.CalcIiwaPosition\n",
    "        )\n",
    "        self.DeclareInitializationDiscreteUpdateEvent(self.Initialize)\n",
    "\n",
    "        self.DeclarePeriodicUnrestrictedUpdateEvent(0.1, 0.0, self.Update)\n",
    "\n",
    "    def Update(self, context, state):\n",
    "        mode = context.get_abstract_state(int(self._mode_index)).get_value()\n",
    "\n",
    "        current_time = context.get_time()\n",
    "        times = context.get_abstract_state(int(self._times_index)).get_value()\n",
    "\n",
    "        if mode == PlannerState.WAIT_FOR_OBJECTS_TO_SETTLE:\n",
    "            if context.get_time() - times[\"initial\"] > 1.0:\n",
    "                self.Plan(context, state)\n",
    "            return\n",
    "        elif mode == PlannerState.GO_HOME:\n",
    "            traj_q = context.get_mutable_abstract_state(\n",
    "                int(self._traj_q_index)\n",
    "            ).get_value()\n",
    "            if not traj_q.is_time_in_range(current_time):\n",
    "                self.Plan(context, state)\n",
    "            return\n",
    "\n",
    "        # If we are between pick and place and the gripper is closed, then\n",
    "        # we've missed or dropped the object.  Time to replan.\n",
    "        if (\n",
    "            current_time > times[\"postpick\"]\n",
    "            and current_time < times[\"preplace\"]\n",
    "        ):\n",
    "            wsg_state = self.get_input_port(self._wsg_state_index).Eval(\n",
    "                context\n",
    "            )\n",
    "            if wsg_state[0] < 0.01:\n",
    "                attempts = state.get_mutable_discrete_state(\n",
    "                    int(self._attempts_index)\n",
    "                ).get_mutable_value()\n",
    "                if attempts[0] > 5:\n",
    "                    # If I've failed 5 times in a row, then switch bins.\n",
    "                    print(\n",
    "                        \"Switching to the other bin after 5 consecutive failed attempts\"\n",
    "                    )\n",
    "                    attempts[0] = 0\n",
    "                    if mode == PlannerState.PICKING_FROM_X_BIN:\n",
    "                        state.get_mutable_abstract_state(\n",
    "                            int(self._mode_index)\n",
    "                        ).set_value(PlannerState.PICKING_FROM_Y_BIN)\n",
    "                    else:\n",
    "                        state.get_mutable_abstract_state(\n",
    "                            int(self._mode_index)\n",
    "                        ).set_value(PlannerState.PICKING_FROM_X_BIN)\n",
    "                    self.Plan(context, state)\n",
    "                    return\n",
    "\n",
    "                attempts[0] += 1\n",
    "                state.get_mutable_abstract_state(\n",
    "                    int(self._mode_index)\n",
    "                ).set_value(PlannerState.WAIT_FOR_OBJECTS_TO_SETTLE)\n",
    "                times = {\"initial\": current_time}\n",
    "                state.get_mutable_abstract_state(\n",
    "                    int(self._times_index)\n",
    "                ).set_value(times)\n",
    "                X_G = self.get_input_port(0).Eval(context)[\n",
    "                    int(self._gripper_body_index)\n",
    "                ]\n",
    "                state.get_mutable_abstract_state(\n",
    "                    int(self._traj_X_G_index)\n",
    "                ).set_value(\n",
    "                    PiecewisePose.MakeLinear(\n",
    "                        [current_time, np.inf], [X_G, X_G]\n",
    "                    )\n",
    "                )\n",
    "                return\n",
    "\n",
    "        traj_X_G = context.get_abstract_state(\n",
    "            int(self._traj_X_G_index)\n",
    "        ).get_value()\n",
    "        if not traj_X_G.is_time_in_range(current_time):\n",
    "            self.Plan(context, state)\n",
    "            return\n",
    "\n",
    "        X_G = self.get_input_port(0).Eval(context)[\n",
    "            int(self._gripper_body_index)\n",
    "        ]\n",
    "        # if current_time > 10 and current_time < 12:\n",
    "        #    self.GoHome(context, state)\n",
    "        #    return\n",
    "        if (\n",
    "            np.linalg.norm(\n",
    "                traj_X_G.GetPose(current_time).translation()\n",
    "                - X_G.translation()\n",
    "            )\n",
    "            > 0.2\n",
    "        ):\n",
    "            # If my trajectory tracking has gone this wrong, then I'd better\n",
    "            # stop and replan.  TODO(russt): Go home, in joint coordinates,\n",
    "            # instead.\n",
    "            self.GoHome(context, state)\n",
    "            return\n",
    "\n",
    "    def GoHome(self, context, state):\n",
    "        print(\"Replanning due to large tracking error.\")\n",
    "        state.get_mutable_abstract_state(int(self._mode_index)).set_value(\n",
    "            PlannerState.GO_HOME\n",
    "        )\n",
    "        q = self.get_input_port(self._iiwa_position_index).Eval(context)\n",
    "        q0 = copy(context.get_discrete_state(self._q0_index).get_value())\n",
    "        q0[0] = q[0]  # Safer to not reset the first joint.\n",
    "\n",
    "        current_time = context.get_time()\n",
    "        q_traj = PiecewisePolynomial.FirstOrderHold(\n",
    "            [current_time, current_time + 5.0], np.vstack((q, q0)).T\n",
    "        )\n",
    "        state.get_mutable_abstract_state(int(self._traj_q_index)).set_value(\n",
    "            q_traj\n",
    "        )\n",
    "\n",
    "    def Plan(self, context, state):\n",
    "        mode = copy(\n",
    "            state.get_mutable_abstract_state(int(self._mode_index)).get_value()\n",
    "        )\n",
    "\n",
    "        X_G = {\n",
    "            \"initial\": self.get_input_port(0).Eval(context)[\n",
    "                int(self._gripper_body_index)\n",
    "            ]\n",
    "        }\n",
    "\n",
    "        cost = np.inf\n",
    "        for i in range(5):\n",
    "            if mode == PlannerState.PICKING_FROM_Y_BIN:\n",
    "                cost, X_G[\"pick\"] = self.get_input_port(\n",
    "                    self._y_bin_grasp_index\n",
    "                ).Eval(context)\n",
    "                if np.isinf(cost):\n",
    "                    mode = PlannerState.PICKING_FROM_X_BIN\n",
    "            else:\n",
    "                cost, X_G[\"pick\"] = self.get_input_port(\n",
    "                    self._x_bin_grasp_index\n",
    "                ).Eval(context)\n",
    "                if np.isinf(cost):\n",
    "                    mode = PlannerState.PICKING_FROM_Y_BIN\n",
    "                else:\n",
    "                    mode = PlannerState.PICKING_FROM_X_BIN\n",
    "\n",
    "            if not np.isinf(cost):\n",
    "                break\n",
    "\n",
    "        assert not np.isinf(\n",
    "            cost\n",
    "        ), \"Could not find a valid grasp in either bin after 5 attempts\"\n",
    "        state.get_mutable_abstract_state(int(self._mode_index)).set_value(mode)\n",
    "\n",
    "        # TODO(russt): The randomness should come in through a random input\n",
    "        # port.\n",
    "        if mode == PlannerState.PICKING_FROM_X_BIN:\n",
    "            # Place in Y bin:\n",
    "            X_G[\"place\"] = RigidTransform(\n",
    "                RollPitchYaw(-np.pi / 2, 0, 0),\n",
    "                [rng.uniform(-0.25, 0.15), rng.uniform(-0.6, -0.4), 0.3],\n",
    "            )\n",
    "        else:\n",
    "            # Place in X bin:\n",
    "            X_G[\"place\"] = RigidTransform(\n",
    "                RollPitchYaw(-np.pi / 2, 0, np.pi / 2),\n",
    "                [rng.uniform(0.35, 0.65), rng.uniform(-0.12, 0.28), 0.3],\n",
    "            )\n",
    "\n",
    "        X_G, times = MakeGripperFrames(X_G, t0=context.get_time())\n",
    "        print(\n",
    "            f\"Planned {times['postplace'] - times['initial']} second trajectory in mode {mode} at time {context.get_time()}.\"\n",
    "        )\n",
    "        state.get_mutable_abstract_state(int(self._times_index)).set_value(\n",
    "            times\n",
    "        )\n",
    "\n",
    "        if False:  # Useful for debugging\n",
    "            AddMeshcatTriad(meshcat, \"X_Oinitial\", X_PT=X_O[\"initial\"])\n",
    "            AddMeshcatTriad(meshcat, \"X_Gprepick\", X_PT=X_G[\"prepick\"])\n",
    "            AddMeshcatTriad(meshcat, \"X_Gpick\", X_PT=X_G[\"pick\"])\n",
    "            AddMeshcatTriad(meshcat, \"X_Gplace\", X_PT=X_G[\"place\"])\n",
    "\n",
    "        traj_X_G = MakeGripperPoseTrajectory(X_G, times)\n",
    "        traj_wsg_command = MakeGripperCommandTrajectory(times)\n",
    "\n",
    "        state.get_mutable_abstract_state(int(self._traj_X_G_index)).set_value(\n",
    "            traj_X_G\n",
    "        )\n",
    "        state.get_mutable_abstract_state(int(self._traj_wsg_index)).set_value(\n",
    "            traj_wsg_command\n",
    "        )\n",
    "\n",
    "    def start_time(self, context):\n",
    "        return (\n",
    "            context.get_abstract_state(int(self._traj_X_G_index))\n",
    "            .get_value()\n",
    "            .start_time()\n",
    "        )\n",
    "\n",
    "    def end_time(self, context):\n",
    "        return (\n",
    "            context.get_abstract_state(int(self._traj_X_G_index))\n",
    "            .get_value()\n",
    "            .end_time()\n",
    "        )\n",
    "\n",
    "    def CalcGripperPose(self, context, output):\n",
    "        context.get_abstract_state(int(self._mode_index)).get_value()\n",
    "\n",
    "        traj_X_G = context.get_abstract_state(\n",
    "            int(self._traj_X_G_index)\n",
    "        ).get_value()\n",
    "        if traj_X_G.get_number_of_segments() > 0 and traj_X_G.is_time_in_range(\n",
    "            context.get_time()\n",
    "        ):\n",
    "            # Evaluate the trajectory at the current time, and write it to the\n",
    "            # output port.\n",
    "            output.set_value(\n",
    "                context.get_abstract_state(int(self._traj_X_G_index))\n",
    "                .get_value()\n",
    "                .GetPose(context.get_time())\n",
    "            )\n",
    "            return\n",
    "\n",
    "        # Command the current position (note: this is not particularly good if the velocity is non-zero)\n",
    "        output.set_value(\n",
    "            self.get_input_port(0).Eval(context)[int(self._gripper_body_index)]\n",
    "        )\n",
    "\n",
    "    def CalcWsgPosition(self, context, output):\n",
    "        mode = context.get_abstract_state(int(self._mode_index)).get_value()\n",
    "        opened = np.array([0.107])\n",
    "        np.array([0.0])\n",
    "\n",
    "        if mode == PlannerState.GO_HOME:\n",
    "            # Command the open position\n",
    "            output.SetFromVector([opened])\n",
    "            return\n",
    "\n",
    "        traj_wsg = context.get_abstract_state(\n",
    "            int(self._traj_wsg_index)\n",
    "        ).get_value()\n",
    "        if traj_wsg.get_number_of_segments() > 0 and traj_wsg.is_time_in_range(\n",
    "            context.get_time()\n",
    "        ):\n",
    "            # Evaluate the trajectory at the current time, and write it to the\n",
    "            # output port.\n",
    "            output.SetFromVector(traj_wsg.value(context.get_time()))\n",
    "            return\n",
    "\n",
    "        # Command the open position\n",
    "        output.SetFromVector([opened])\n",
    "\n",
    "    def CalcControlMode(self, context, output):\n",
    "        mode = context.get_abstract_state(int(self._mode_index)).get_value()\n",
    "\n",
    "        if mode == PlannerState.GO_HOME:\n",
    "            output.set_value(InputPortIndex(2))  # Go Home\n",
    "        else:\n",
    "            output.set_value(InputPortIndex(1))  # Diff IK\n",
    "\n",
    "    def CalcDiffIKReset(self, context, output):\n",
    "        mode = context.get_abstract_state(int(self._mode_index)).get_value()\n",
    "\n",
    "        if mode == PlannerState.GO_HOME:\n",
    "            output.set_value(True)\n",
    "        else:\n",
    "            output.set_value(False)\n",
    "\n",
    "    def Initialize(self, context, discrete_state):\n",
    "        discrete_state.set_value(\n",
    "            int(self._q0_index),\n",
    "            self.get_input_port(int(self._iiwa_position_index)).Eval(context),\n",
    "        )\n",
    "\n",
    "    def CalcIiwaPosition(self, context, output):\n",
    "        traj_q = context.get_mutable_abstract_state(\n",
    "            int(self._traj_q_index)\n",
    "        ).get_value()\n",
    "\n",
    "        output.SetFromVector(traj_q.value(context.get_time()))\n",
    "\n",
    "\n",
    "def clutter_clearing_demo():\n",
    "    meshcat.Delete()\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    scenario = load_scenario(\n",
    "        filename=FindResource(\"models/clutter.scenarios.yaml\"),\n",
    "        scenario_name=\"Clutter\",\n",
    "    )\n",
    "    model_directives = \"\"\"\n",
    "directives:\n",
    "\"\"\"\n",
    "    for i in range(10 if running_as_notebook else 2):\n",
    "        object_num = rng.integers(0, len(ycb))\n",
    "        if \"cracker_box\" in ycb[object_num]:\n",
    "            # skip it. it's just too big!\n",
    "            continue\n",
    "        model_directives += f\"\"\"\n",
    "- add_model:\n",
    "    name: ycb{i}\n",
    "    file: package://manipulation/hydro/{ycb[object_num]}\n",
    "\"\"\"\n",
    "    scenario = add_directives(scenario, data=model_directives)\n",
    "\n",
    "    station = builder.AddSystem(MakeHardwareStation(scenario, meshcat))\n",
    "    to_point_cloud = AddPointClouds(\n",
    "        scenario=scenario, station=station, builder=builder\n",
    "    )\n",
    "    plant = station.GetSubsystemByName(\"plant\")\n",
    "\n",
    "    y_bin_grasp_selector = builder.AddSystem(\n",
    "        GraspSelector(\n",
    "            plant,\n",
    "            plant.GetModelInstanceByName(\"bin0\"),\n",
    "            camera_body_indices=[\n",
    "                plant.GetBodyIndices(plant.GetModelInstanceByName(\"camera0\"))[\n",
    "                    0\n",
    "                ],\n",
    "                plant.GetBodyIndices(plant.GetModelInstanceByName(\"camera1\"))[\n",
    "                    0\n",
    "                ],\n",
    "                plant.GetBodyIndices(plant.GetModelInstanceByName(\"camera2\"))[\n",
    "                    0\n",
    "                ],\n",
    "            ],\n",
    "        )\n",
    "    )\n",
    "    builder.Connect(\n",
    "        to_point_cloud[\"camera0\"].get_output_port(),\n",
    "        y_bin_grasp_selector.get_input_port(0),\n",
    "    )\n",
    "    builder.Connect(\n",
    "        to_point_cloud[\"camera1\"].get_output_port(),\n",
    "        y_bin_grasp_selector.get_input_port(1),\n",
    "    )\n",
    "    builder.Connect(\n",
    "        to_point_cloud[\"camera2\"].get_output_port(),\n",
    "        y_bin_grasp_selector.get_input_port(2),\n",
    "    )\n",
    "    builder.Connect(\n",
    "        station.GetOutputPort(\"body_poses\"),\n",
    "        y_bin_grasp_selector.GetInputPort(\"body_poses\"),\n",
    "    )\n",
    "\n",
    "    x_bin_grasp_selector = builder.AddSystem(\n",
    "        GraspSelector(\n",
    "            plant,\n",
    "            plant.GetModelInstanceByName(\"bin1\"),\n",
    "            camera_body_indices=[\n",
    "                plant.GetBodyIndices(plant.GetModelInstanceByName(\"camera3\"))[\n",
    "                    0\n",
    "                ],\n",
    "                plant.GetBodyIndices(plant.GetModelInstanceByName(\"camera4\"))[\n",
    "                    0\n",
    "                ],\n",
    "                plant.GetBodyIndices(plant.GetModelInstanceByName(\"camera5\"))[\n",
    "                    0\n",
    "                ],\n",
    "            ],\n",
    "        )\n",
    "    )\n",
    "    builder.Connect(\n",
    "        to_point_cloud[\"camera3\"].get_output_port(),\n",
    "        x_bin_grasp_selector.get_input_port(0),\n",
    "    )\n",
    "    builder.Connect(\n",
    "        to_point_cloud[\"camera4\"].get_output_port(),\n",
    "        x_bin_grasp_selector.get_input_port(1),\n",
    "    )\n",
    "    builder.Connect(\n",
    "        to_point_cloud[\"camera5\"].get_output_port(),\n",
    "        x_bin_grasp_selector.get_input_port(2),\n",
    "    )\n",
    "    builder.Connect(\n",
    "        station.GetOutputPort(\"body_poses\"),\n",
    "        x_bin_grasp_selector.GetInputPort(\"body_poses\"),\n",
    "    )\n",
    "\n",
    "    planner = builder.AddSystem(Planner(plant))\n",
    "    builder.Connect(\n",
    "        station.GetOutputPort(\"body_poses\"), planner.GetInputPort(\"body_poses\")\n",
    "    )\n",
    "    builder.Connect(\n",
    "        x_bin_grasp_selector.get_output_port(),\n",
    "        planner.GetInputPort(\"x_bin_grasp\"),\n",
    "    )\n",
    "    builder.Connect(\n",
    "        y_bin_grasp_selector.get_output_port(),\n",
    "        planner.GetInputPort(\"y_bin_grasp\"),\n",
    "    )\n",
    "    builder.Connect(\n",
    "        station.GetOutputPort(\"wsg.state_measured\"),\n",
    "        planner.GetInputPort(\"wsg_state\"),\n",
    "    )\n",
    "    builder.Connect(\n",
    "        station.GetOutputPort(\"iiwa.position_measured\"),\n",
    "        planner.GetInputPort(\"iiwa_position\"),\n",
    "    )\n",
    "\n",
    "    robot = station.GetSubsystemByName(\n",
    "        \"iiwa.controller\"\n",
    "    ).get_multibody_plant_for_control()\n",
    "\n",
    "    # Set up differential inverse kinematics.\n",
    "    diff_ik = AddIiwaDifferentialIK(builder, robot)\n",
    "    builder.Connect(planner.GetOutputPort(\"X_WG\"), diff_ik.get_input_port(0))\n",
    "    builder.Connect(\n",
    "        station.GetOutputPort(\"iiwa.state_estimated\"),\n",
    "        diff_ik.GetInputPort(\"robot_state\"),\n",
    "    )\n",
    "    builder.Connect(\n",
    "        planner.GetOutputPort(\"reset_diff_ik\"),\n",
    "        diff_ik.GetInputPort(\"use_robot_state\"),\n",
    "    )\n",
    "\n",
    "    builder.Connect(\n",
    "        planner.GetOutputPort(\"wsg_position\"),\n",
    "        station.GetInputPort(\"wsg.position\"),\n",
    "    )\n",
    "\n",
    "    # The DiffIK and the direct position-control modes go through a PortSwitch\n",
    "    switch = builder.AddSystem(PortSwitch(7))\n",
    "    builder.Connect(\n",
    "        diff_ik.get_output_port(), switch.DeclareInputPort(\"diff_ik\")\n",
    "    )\n",
    "    builder.Connect(\n",
    "        planner.GetOutputPort(\"iiwa_position_command\"),\n",
    "        switch.DeclareInputPort(\"position\"),\n",
    "    )\n",
    "    builder.Connect(\n",
    "        switch.get_output_port(), station.GetInputPort(\"iiwa.position\")\n",
    "    )\n",
    "    builder.Connect(\n",
    "        planner.GetOutputPort(\"control_mode\"),\n",
    "        switch.get_port_selector_input_port(),\n",
    "    )\n",
    "\n",
    "    visualizer = MeshcatVisualizer.AddToBuilder(\n",
    "        builder, station.GetOutputPort(\"query_object\"), meshcat\n",
    "    )\n",
    "    diagram = builder.Build()\n",
    "\n",
    "    simulator = Simulator(diagram)\n",
    "    context = simulator.get_context()\n",
    "\n",
    "    plant_context = plant.GetMyMutableContextFromRoot(context)\n",
    "    z = 0.2\n",
    "    for body_index in plant.GetFloatingBaseBodies():\n",
    "        tf = RigidTransform(\n",
    "            UniformlyRandomRotationMatrix(generator),\n",
    "            [rng.uniform(0.35, 0.65), rng.uniform(-0.12, 0.28), z],\n",
    "        )\n",
    "        plant.SetFreeBodyPose(plant_context, plant.get_body(body_index), tf)\n",
    "        z += 0.1\n",
    "\n",
    "    simulator.AdvanceTo(0.1)\n",
    "    meshcat.Flush()  # Wait for the large object meshes to get to meshcat.\n",
    "\n",
    "    if running_as_notebook:\n",
    "        simulator.set_target_realtime_rate(1.0)\n",
    "        meshcat.AddButton(\"Stop Simulation\", \"Escape\")\n",
    "        print(\"Press Escape to stop the simulation\")\n",
    "        while meshcat.GetButtonClicks(\"Stop Simulation\") < 1:\n",
    "            simulator.AdvanceTo(simulator.get_context().get_time() + 2.0)\n",
    "        meshcat.DeleteButton(\"Stop Simulation\")\n",
    "\n",
    "\n",
    "clutter_clearing_demo()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
